from datetime import datetime
import streamlit as st
from data.user_data import obtain_data, modify_data

# https://blog.csdn.net/yuan2019035055/article/details/129754864
st.set_page_config(
    page_title="goudaner.AI",
    page_icon="🧊",
    layout="wide",
    initial_sidebar_state="expanded"
)
from chat_agent import get_llm, obtain_answer
from streamlit_option_menu import option_menu
from data.user_data import insert_data
import os
from customer_logging import get_logger

os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"
from vanna_calls import (
    generate_sql_cached,
    run_sql_cached,
    generate_plotly_code_cached,
    generate_plot_cached,
    get_vanna,
)

logger = get_logger("sys-")


def upload_Save():
    if "uploaded_state" not in st.session_state:
        st.session_state.uploaded_state = False
    if st.session_state.uploaded_state:
        return
    # 创建一个文件夹用于保存上传的文件（若存在则清空，若不存在，则新建）
    dirs = 'One'
    if not os.path.exists(dirs):
        os.makedirs(dirs)
    # 选择文件
    uploaded_files = st.file_uploader("###### 您可以再此处上传知识文件", accept_multiple_files=True,
                                      type=["pdf", "txt", "docx", "jpg", "png"])
    # 保存文件
    if uploaded_files is None:
        return
    for uploaded_file in uploaded_files:
        file_contents = uploaded_file.getvalue()
        file_path = os.path.join(dirs, uploaded_file.name)

        llm = get_llm()
        fassStore = llm.store
        # 将文件保存到本地文件系统
        with open(file_path, "wb") as f:
            f.write(file_contents)
        # 获取文件路径
        fassStore.load_file(file_path)
        st.session_state.uploaded_state = True
        st.write(f"文件地址: {file_path}")
    return os.path.join(os.path.dirname(os.path.abspath(__file__)), dirs)


def_question = """🚀 狗蛋1号 问题示例：订单20230331064371000533的回寄地址是多少？                    
                            订单20230331064371000533的换货单号?     
                            吉利端午礼盒 包含什么？"""
def_data_question = """🚀 狗蛋1号 问题示例：2024-03 的各分类商品 销售额 top10商品 排除掉 已退款订单                                                 
                                吉好的商城2024年4月 每天几点钟下单人数是多少 已小时为间隔统计                                                                  
吉好的商城，spuid=15845的商品，销售总金额为多少？其中吉点点总金额为多少？现金多少？其他方式销售金额多少？优惠金额为多少？总共卖了多少坛？"""
def_message = "尊敬的用户您好，您的吉利伙伴assistant很高兴为您服务，请问有什么可以帮您？"


def main():
    st.sidebar.title("goudaner.AI v0.1")
    menu1 = "知识库ai"
    menu4 = "添加-知识"
    menu2 = "数据统筹功能"
    menu3 = "数据统筹训练数据"
    with st.sidebar:
        menu = option_menu("功能分类", [menu1, menu4, menu2, menu3],
                           icons=['house', "list-task"],
                           menu_icon="cast", default_index=0)
    get_vanna()
    get_llm()
    obtain_answer()
    if menu == menu1:
        st.title("💬 知识库ai")
        st.caption(def_question)
        if st.button('清除聊天记录'):
            st.session_state.history = [{"role": "assistant", "content": def_message}]
        showLLMChatbot()
    if menu == menu4:
        user_training_data = obtain_data("""SELECT uid,type,content FROM `user_training` where del=0""")
        # 字典
        st.dataframe(
            data=user_training_data
        )
        # 文件上传
        upload_Save()
        knowledge_training_data()
    if menu == menu2:
        st.subheader("数据列表")
        st.caption(def_data_question)
        showdataframe()

    if menu == menu3:
        st.subheader("数据统筹训练数据")
        training_data()


def knowledge_training_data():
    llm = get_llm()
    faiss = llm.store
    with st.form("del_calc"):
        submitted = st.form_submit_button("删除")
        data_id = st.text_input("data_id")
        # 检查是否提交了表单
        if submitted:
            try:
                faiss.del_data([data_id])
            except ValueError as e:
                logger.error(f"删除不存在uid={data_id}")
            modify_data(f"""UPDATE user_training SET del = 1 WHERE uid = \"{data_id}\"""")
            st.experimental_rerun()
    # 定义selectbox来获取数据类型
    data_type = st.selectbox("训练数据类型", ("商品信息", "问答示例"))
    # 使用st.form来处理提交事件
    with st.form("t_d"):
        training_type = data_type == "商品信息"
        input_type = "商品信息" if training_type else "问题"
        answer_type = "商品详情" if training_type else "回答"
        input_info = st.text_input(input_type, max_chars=100, help='最大长度为100字符')
        answer_info = st.text_area(label=answer_type,
                                   value='请输入...',
                                   height=5,
                                   max_chars=300,
                                   help='最大长度限制为300')
        # 提交按钮
        submitted = st.form_submit_button("训练")
        if submitted:

            if input_type == "商品信息":
                template = f"""Product information : {input_info} ;   {answer_info} """
            else:
                template = f"""Q&A examples : {input_info} ;   {answer_info} """
            uid = faiss.write_templates(template)

            user_training_data = {
                "uid": uid,
                "type": input_type,
                "content": template,
                "state": "Untrained"
            }
            insert_data("user_training", [user_training_data])
            st.experimental_rerun()


def training_data():
    vn = get_vanna()
    df = vn.get_training_data()

    if df is None or len(df) == 0:
        st.text("暂无训练数据")
        return
    # 字典
    st.dataframe(
        data=df
    )

    # 定义selectbox来获取数据类型
    data_type = st.selectbox("训练数据类型", ("DDL", "Documentation", "SQL"))

    # 使用st.form来处理提交事件
    with st.form("calc"):
        question = ''
        documentation = ''

        # 根据数据类型显示不同的输入字段
        if data_type == "SQL":
            question = st.text_input("问题")
            sql = st.text_input("SQL")
        else:
            documentation = st.text_area("请在此添加训练数据", "如：DDL，Doc，SQL", height=200)

            # 提交按钮
        submitted = st.form_submit_button("训练")
        # 检查是否提交了表单
        if submitted:
            # 根据数据类型调用train函数，并传递相应的参数
            if data_type == "SQL":
                vn.train(question=question, sql=sql)
            else:
                vn.train(ddl=documentation, documentation=documentation)

                # 显示结果或进行其他操作
            # st.write(f"已经添加了 {data_type} 数据，实际数据：{question or documentation}")

            "这里就可以添加数据了" + data_type + ",实际数据：" + question
            st.experimental_rerun()

    with st.form("del_calc"):
        submitted = st.form_submit_button("删除")
        data_id = st.text_input("data_id")
        # 检查是否提交了表单
        if submitted:
            vn.remove_training_data(data_id)
            st.experimental_rerun()


# 显示dataframe
def showdataframe():
    # 给对话增加history属性，将历史对话信息储存下来
    if "data_history" not in st.session_state:
        st.session_state.data_history = []
    if "my_question" not in st.session_state:
        st.session_state.my_question = None
    if "sql" not in st.session_state:
        st.session_state.sql = None
    # 显示历史信息
    input_question = st.chat_input()
    my_question = st.session_state.my_question
    if my_question:
        if input_question:
            my_question = my_question + "," + input_question
    else:
        my_question = input_question
    if my_question is None:
        return
    chart_redering(my_question)


@st.cache_data(show_spinner="正在努力生成图表 ...", ttl=60 * 60, experimental_allow_widgets=True)
def chart_redering(my_question):
    chart_logger = get_logger("chart_redering-")
    chart_logger.info("用户问题：" + my_question)
    # 将用户的输入加入历史
    st.session_state.data_history.append({"role": "user", "content": my_question})
    st.session_state.my_question = my_question
    # 在页面上显示用户的输入
    with st.chat_message("user"):
        st.markdown(my_question)

    sql = st.session_state.sql
    if sql is None:
        sql = generate_sql_cached(my_question)
    if sql is None:
        assistant_message_error = st.chat_message(
            "assistant", avatar="https://i.geely.com/favicon.ico"
        )
        assistant_message_error.error("我无法为该问题生成SQL")
        return
    chart_logger.info("用户问题生成SQL：" + sql)
    assistant_message_sql = st.chat_message(
        "assistant", avatar="https://i.geely.com/favicon.ico"
    )
    assistant_message_sql.code(sql, language="sql", line_numbers=True)
    # 将模型的输出加入到历史信息中
    st.session_state.data_history.append({"role": "assistant", "content": sql})
    df = None
    user_message_sql_check = st.chat_message("assistant", avatar="https://i.geely.com/favicon.ico")
    user_message_sql_check.write(f"您对生成的SQL代码满意吗？")
    with user_message_sql_check:
        happy_sql = st.radio(
            "Happy",
            options=["", "满意", "不满意"],
            key="radio_sql",
            index=0,
        )
    if happy_sql == "":
        return
    if happy_sql == "不满意":
        with st.form("sql_change"):
            user_message_sql_change = st.chat_message("assistant", avatar="https://i.geely.com/favicon.ico")
            fixed_sql_query = user_message_sql_change.text_area("请修改sql", sql, height=200)
            # 提交按钮
            submitted = st.form_submit_button("生成")
            # 检查是否提交了表单
            if submitted:
                df = run_sql_cached(sql=fixed_sql_query)
    if happy_sql == "满意":
        # 画一个查询表格到页面
        df = run_sql_cached(sql)
    if df is None:
        return
    st.dataframe(df, use_container_width=True)
    # 生成图标渲染代码
    code = generate_plotly_code_cached(question=my_question, sql=sql, df=df)
    fig = generate_plot_cached(code=code, df=df)
    assistant_message_chart = st.chat_message(
        "assistant", avatar="https://i.geely.com/favicon.ico"
    )
    assistant_message_chart.plotly_chart(fig, use_container_width=True)

    user_message_fig = st.chat_message("assistant", avatar="https://i.geely.com/favicon.ico")
    user_message_fig.write("图表已生成，如需修改，请追加需求")


@st.cache_data(show_spinner="正在努力生成答案 ...", ttl=60 * 60)
def get_response_material(user_input, history):
    print("------------------------------------------------------------")
    logger.info("new用户对话问题：" + user_input)
    agent_executor = obtain_answer()
    response = agent_executor.invoke({"input": user_input, "chat_history": history})["output"]
    logger.info("new用户对话回答：" + response)
    # 返回两个值
    return response


def showLLMChatbot():
    # 给对话增加history属性，将历史对话信息储存下来
    if "history" not in st.session_state:
        st.session_state.history = []
        assistant_message = st.chat_message(
            "assistant", avatar="https://i.geely.com/favicon.ico"
        )
        assistant_message.markdown(def_message)
    # 显示历史信息
    for message in st.session_state.history:
        with st.chat_message(message["role"],
                             avatar=("https://i.geely.com/favicon.ico" if message["role"] == "assistant" else None)):
            st.markdown(message["content"])
    user_input = st.chat_input()
    if user_input is None:
        return
    # 在页面上显示用户的输入
    with st.chat_message("user"):
        st.markdown(user_input)

    response = get_response_material(user_input, st.session_state.history)
    # 在页面上显示模型生成的回复
    with st.chat_message("assistant", avatar="https://i.geely.com/favicon.ico"):
        st.markdown(response)
    # 将用户的输入加入历史
    st.session_state.history.append({"role": "user", "content": user_input})
    # 将模型的输出加入到历史信息中
    st.session_state.history.append({"role": "assistant", "content": response})
    # 只保留十轮对话，这个可根据自己的情况设定，我这里主要是会把history给大模型，context有限，轮数不能太多
    if len(st.session_state.history) > 4:
        st.session_state.history = st.session_state.history[-4:]


if __name__ == '__main__':
    main()
